# 自定义数据集
import os.path as osp
from typing import Union, List, Tuple

import torch
from torch_geometric.data import Dataset, InMemoryDataset, download_url, Data

"""
PyG将自建数据集分为两个文件夹---raw_dir、processed_dir。row_dir是原始的数据集，processed_dir是PyG处理之后的数据集

对于数据集PyG有三种过滤方法---transform、pre_transform、pre_filter。
transform:读取数据，然后对其进行变换
pre_transform：对于整个数据集进行变换，然后将变换之后的数据进行存储，pre_filter同理

PyG将数据集分为两种类型
torch_geometric.data.InMemoryDataset：能够完全放入内存中的
torch_geometric.data.Dataset:不能够完全放入内存中的
"""

"""
创建一个能够完全放入内存中的图数据集InMemoryDataset，但通常创建的数据集不能够放入内存
步骤：
    1、实现torch_geometric.data.InMemoryDataset.raw_file_names():告诉PyG数据集放在哪里
    2、实现torch_geometric.data.InMemoryDataset.processed_file_names():告诉PyG数据集处理完之后放在哪里
    3、实现torch_geometric.data.InMemoryDataset.download():告诉PyG从哪里获取数据集
    4、实现torch_geometric.data.InMemoryDataset.process():告诉PyG如何处理你的数据集
"""

url = ""


# 实现In Memory Dataset的通用模板
class MyDataset(InMemoryDataset):
    # 初始化
    def __init__(self, root, transfrom=None, pre_transform=None):
        # root是数据集的根目录
        super(MyDataset, self).__init__(root, transfrom, pre_transform)
        # 加载数据集
        self.data, self.slices = torch.load(self.processed_paths[0])

    def raw_file_names(self) -> Union[str, List[str], Tuple]:
        return ['file_1', 'file_2', ...]

    def processed_file_names(self) -> Union[str, List[str], Tuple]:
        return ['data.pt']

    def download(self):
        # 将数据集下载到raw_dir文件夹中
        download_url(url, self.raw_dir)

    def process(self):
        data_list = [...]
        # 进行数据过滤
        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]
        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]
        # self.collate将所有数据组合在一起,加速存储
        # data是组合之后的数据
        # slices是分割方式，告诉PyG如何将data还原为原先的数据
        data, slices = self.collate(data_list)
        # 保存数据
        torch.save((data, slices), self.processed_paths[0])


"""
创建一个不能够一次性放入内存的数据集Dataset，这种就类似于PyTorch中的Dataset了，在创建InMemoryDataset的基础上还需要以下步骤
步骤：
    1、实现torch_geometric.data.Dataset.len():告诉PyG数据集有多大
    2、实现torch_geometric.data.Dataset.get():告诉PyG如何从数据集中获取一个数据
"""


# 通用模板如下：
class MyDataset(Dataset):
    # 初始化
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyDataset, self).__init__(root, transform, pre_transform)

    def raw_file_names(self) -> Union[str, List[str], Tuple]:
        return ['file_1', 'file_2', ...]

    def processed_file_names(self) -> Union[str, List[str], Tuple]:
        return ['data_1.pt', ...]

    def download(self):
        path = download_url(url, self.raw_dir)

    def process(self):
        i = 0
        for raw_path in self.raw_paths:
            # 读取数据
            data = Data(...)
            # 过滤数据集
            if self.pre_filter is not None and not self.pre_filter(data):
                pass
            if self.pre_transform is not None:
                data = self.pre_transform(data)
            # 保存数据
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self, idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data
